#include "ocrorast.h"


param_int debug_xycut_intermediate("debug",0,"print intermediate results"
        "to stdout");

// compute integral image for BINARY images
static void compute_integral_image(bytearray& img, 
        intarray& int_img_hor, 
        intarray& int_img_ver){
    int_img_ver.resize(img.dim(0), img.dim(1)) ;
    int_img_hor.resize(img.dim(0), img.dim(1)) ;
    for (int x = 0; x < int_img_ver.dim(0); x++) {
        for (int y = 0; y < int_img_ver.dim(1); y++) {
            int_img_ver(x,y) = 0 ;
            int_img_hor(x,y) = 0 ;
        }
    }
    for (int x = 0; x < int_img_ver.dim(0); x++) {
        int count = 0 ;
        for (int y = 0; y < int_img_ver.dim(1); y++) {
            if (img(x,y) == 0) {
                count++ ;
            }
            int_img_ver(x,y) = count ;
        }
    }

    for (int y = 0; y < int_img_hor.dim(1); y++) {
        int count = 0 ;
        for (int x = 0; x < int_img_hor.dim(0); x++) {
            if (img(x,y) == 0) {
                count++ ;
            }
            int_img_hor(x,y) = count ;
        }
    }
}

// Compute horizontal and vertical projection profiles for a given rectangle
// r out of integral image
static void get_projection_profiles(intarray& proj_on_yaxis, 
        intarray& proj_on_xaxis,
        rectangle& r, 
        intarray& int_img_hor, 
        intarray& int_img_ver) {
    // Mao. Fig.3. 2.a)
    proj_on_yaxis.resize((int)r.y1 - (int)r.y0) ;
    fill(proj_on_yaxis,0);

    proj_on_xaxis.resize((int)r.x1 - (int)r.x0) ;
    fill(proj_on_xaxis,0);

    // JvB: y<= (int)...: all pixels inside the rectangle have to be included.
    for (int y = (int)r.y0; y < (int)r.y1; y++){
        int m=0;
        if(r.x0-1>=0)
            m=int_img_hor((int)r.x0-1, y);
        proj_on_yaxis(y - (int)r.y0) = int_img_hor((int)r.x1-1, y) - m;
    }
    for (int x = (int)r.x0; x < (int)r.x1; x++){
        int m=0;
        if(r.y0-1>=0)
            m=int_img_ver(x,(int)r.y0-1);
        proj_on_xaxis(x - (int)r.x0) = int_img_ver(x, (int)r.y1-1) - m;
    }
}


// shrink to the bounding boy of all the black pixels in the rectangle and
// remove noise
static void shrink_and_clean(rectangle& r, 
        intarray& proj_on_yaxis, 
        intarray& proj_on_xaxis, 
        double f_tnx, 
        double f_tny) {
    // Mao. Fig.3. 2.b)
    // shrinking
    ASSERT(r.area());
    ASSERT(proj_on_xaxis.dim(0));
    ASSERT(proj_on_xaxis.dim(0));
    ASSERT(sum(proj_on_xaxis)!=0);
    ASSERT(sum(proj_on_xaxis)==sum(proj_on_yaxis));
    int xbegin = 0, xend = proj_on_xaxis.dim(0) - 1 ;
    while (xbegin < proj_on_xaxis.length() && proj_on_xaxis(xbegin) <= 0)
        xbegin++ ;
    while (xend >= 0 && proj_on_xaxis(xend) <= 0)
        xend-- ;
    int ybegin = 0, yend = proj_on_yaxis.dim(0) - 1 ;
    while (ybegin < proj_on_yaxis.length() && proj_on_yaxis(ybegin) <= 0)
        ybegin++ ;
    while (yend >= 0 && proj_on_yaxis(yend) <= 0)
        yend-- ;


    int old_x0 = (int)r.x0 ;
    int old_y0 = (int)r.y0 ;
    r.x0 = old_x0 + xbegin ;
    r.x1 = old_x0 + xend+1 ;
    r.y0 = old_y0 + ybegin ;
    r.y1 = old_y0 + yend+1 ;

    // Mao. Fig.3. 2.b) & 2.c)
    // cleaning
    for (int i = 0; i < proj_on_yaxis.dim(0); i++)
        proj_on_yaxis(i) = proj_on_yaxis(i) - (int)(r.width()*f_tny) ;
    for (int i = 0; i < proj_on_xaxis.dim(0); i++)
        proj_on_xaxis(i) = proj_on_xaxis(i) - (int)(r.height()*f_tnx) ;

}


// returns the widest non-starting or ending white gape of the two
// projection profiles
static void get_widest_gap(int& cut_pos_y, 
        int& gap_y,
        int& cut_pos_x, 
        int& gap_x,  
        intarray& proj_on_yaxis,
        intarray& proj_on_xaxis) {
    int gap_hor = -1, gap_ver = -1, pos_hor = -1, pos_ver = -1 ;
    int begin = -1 ;
    int end = -1 ;
    // find gap in y-axis projection
    for (int i = 1; i < proj_on_yaxis.dim(0); i++) {
        if (begin>=0
                && proj_on_yaxis(i-1) <= 0 
                && proj_on_yaxis(i) > 0)
            end = i ;
        if (proj_on_yaxis(i-1) > 0 && proj_on_yaxis(i) <= 0)
            begin = i ;
        if (begin > 0 && end > 0 && end-begin > gap_hor) {
            gap_hor = end - begin ;
            pos_hor = (end + begin) / 2 ;
            begin = -1 ;
            end = -1 ;
        }
    }

    begin = -1 ;
    end = -1 ;
    // find gap in x-axis projection
    for (int i = 1; i < proj_on_xaxis.dim(0); i++) {
        if (begin>=0
                && proj_on_xaxis(i-1) <= 0 
                && proj_on_xaxis(i) > 0)
            end = i ;
        if (proj_on_xaxis(i-1) > 0 && proj_on_xaxis(i) <= 0)
            begin = i ;
        if (begin > 0 && end > 0 && end-begin > gap_ver) {
            gap_ver = end - begin ;
            pos_ver = (end + begin) / 2 ;
            begin = -1 ;
            end = -1 ;
        }
    }

    cut_pos_y = pos_hor;
    gap_y = gap_hor;

    cut_pos_x = pos_ver;
    gap_x = gap_ver;

}

// Split rectangle by returning the coordinates of two new rectangles r1, r2
static void split_rect(rectangle& r1, 
        rectangle& r2, 
        rectangle& r, 
        int dir, 
        int pos) {
    if (dir == HORIZONTAL_CUT) { // horizontal cut
        r1 = rectangle(r.x0, r.y0, r.x1, pos) ; // lower rectangle
        r2 = rectangle(r.x0, pos, r.x1, r.y1) ; // upper rectangle
    }
    if (dir == VERTICAL_CUT) { // vertical cut
        r1 = rectangle(r.x0, r.y0, pos, r.y1) ; // left rectangle
        r2 = rectangle(pos, r.y0, r.x1, r.y1) ; // rigth rectangle
    }
}

static void insertRectToTree(int index, 
        rectarray& tree, 
        rectangle rL, 
        rectangle rR) {
    int left = index * 2 + 1 ;
    int right = index * 2 + 2 ;
    while (right+1 > tree.length()) // increased size of tree if necessary
        tree.push(rectangle(-1, -1, -1, -1)) ;
    tree[left] = rL ;
    tree[right] = rR ;
}

static void xycut(rectarray& blocks, 
        rectarray& tree, 
        intarray& int_img_hor,
        intarray& int_img_ver, 
        int tnx, int tny, int tcx, int tcy) {
    // arrays to save the projection profile
    intarray proj_on_yaxis ; // his_x
    intarray proj_on_xaxis ; // his_y

    // Mao. Fig.3. 2.0)
    // adding whole page as root node
    rectangle page = rectangle(0, 0, int_img_ver.dim(0),
            int_img_ver.dim(1)) ;
    tree.push(page) ; // push root rectangle to tree

    // compute factor that is needed for noise removal
    if (int_img_hor.dim(0) <= 0 || 
            int_img_ver.dim(1) <= 0 || 
            int_img_hor.length1d() != int_img_ver.length1d()) {
        fprintf(stderr, "ocr-pageseg-xycut: xycut: Error in image dim!\n");
        return;
        //exit(0) ;
    }
    double factor_tnx = (double)tnx / (double)int_img_hor.dim(1) ;
    double factor_tny = (double)tny / (double)int_img_ver.dim(0) ;

    int i = 0 ;
    rectangle r = tree.at(0) ;
    while (i < tree.length()) {
        // take node in queue and process it

        if (tree[i].x0 >= 0 
                && tree[i].y0 >= 0 
                && tree[i].x1 >= 0 
                && tree[i].y1 >= 0) {
            // compute projection profile
            get_projection_profiles(proj_on_yaxis, proj_on_xaxis, 
                    r, int_img_hor, int_img_ver) ;


            // Mao. Fig.3. 2.b) & 2.c)
            // save old starting points of rectangle, before shrinking
            int old_x0 = (int)r.x0 ;
            int old_y0 = (int)r.y0 ;
            // remove noise and shrink to effectively used area
            shrink_and_clean(r, proj_on_yaxis, proj_on_xaxis, factor_tnx, factor_tny);
            // changes the values of r!!!

            // Mao. Fig.3. 2.d)
            //                 int dir = -1 ;
            int cut_pos_y = -1 ;
            int gap_y = -1 ;
            int cut_pos_x = -1 ;
            int gap_x = -1 ;


            // find widest gap and its middle and the direction of the cut to be
            // done (0 = hor, 1 = ver)
            get_widest_gap(cut_pos_y, gap_y, 
                    cut_pos_x, gap_x, 
                    proj_on_yaxis, proj_on_xaxis) ;

            // split in the right direction
            if (gap_y >= gap_x && gap_y > tcy) {
                rectangle rB, rT ; // rectangleTop rectangleBottom
                split_rect(rB, rT, r, HORIZONTAL_CUT, old_y0 + cut_pos_y) ;
                insertRectToTree(i, tree, rB, rT) ;
            }
            else if (gap_y >= gap_x && gap_y <= tcy && gap_x > tcx) {
                rectangle rL, rR ; // rectangleTop rectangleBottom
                split_rect(rL, rR, r, VERTICAL_CUT, old_x0 + cut_pos_x) ;
                insertRectToTree(i, tree, rL, rR) ;
            }
            else if (gap_x > gap_y && gap_x > tcx) {
                rectangle rL, rR ; // rectangleTop rectangleBottom
                split_rect(rL, rR, r, VERTICAL_CUT, old_x0 + cut_pos_x) ;
                insertRectToTree(i, tree, rL, rR) ;
            }
            else if (gap_x > gap_y && gap_x <= tcx && gap_y > tcy) {
                rectangle rL, rR ; // rectangleTop rectangleBottom
                split_rect(rL, rR, r, HORIZONTAL_CUT, old_y0 + cut_pos_y) ;
                insertRectToTree(i, tree, rL, rR) ;
            }

            else {
                //insert dummy rects
                rectangle rL = rectangle(-1, -1, -1, -1) ;
                rectangle rR = rectangle(-1, -1, -1, -1) ;
                // insert leaf rect to result blocks
                rectangle ret_rect = rectangle(r.x0, r.y0, r.x1, r.y1) ;
                blocks.push(ret_rect) ;
            }
        }
        i++ ;
        if(i<tree.length())
            r = tree[i] ;
    }
}


void SegmentPageByXYCUTS::setParameters(unsigned int itnx, 
        unsigned int itny, 
        unsigned int itcx, 
        unsigned int itcy) {
    tnx=itnx ;
    tny=itny ;
    tcx=itcx ;
    tcy=itcy ;
    if (tcx <= 0) {
        fprintf(stderr, "ocr-pageseg-xycut:tcx value too low; tcx = 1\n");
        tcx = 1 ;
    }
    if (tcy <= 0) {
        fprintf(stderr, "ocr-pageseg-xycut:tcy value too low; tcy = 1\n");
        tcy = 1 ;
    }
}

SegmentPageByXYCUTS::SegmentPageByXYCUTS(unsigned int itnx, 
        unsigned int itny, 
        unsigned int itcx, 
        unsigned int itcy) {
    tnx=itnx ;
    tny=itny ;
    tcx=itcx ;
    tcy=itcy ;
    if (tcx <= 0) {
        fprintf(stderr, "ocr-pageseg-xycut:tcx value too low; tcx = 1\n");
        tcx = 1 ;
    }
    if (tcy <= 0) {
        fprintf(stderr, "ocr-pageseg-xycut:tcy value too low; tcy = 1\n");
        tcy = 1 ;
    }
}

SegmentPageByXYCUTS::SegmentPageByXYCUTS() {
    tnx=78 ;
    tny=32 ;
    tcx=35 ;
    tcy=54 ;
}


void SegmentPageByXYCUTS::segment(intarray &image,bytearray &in) {
    if(!contains_only(in,byte(0),byte(255))){
        fprintf(stderr,"X-Y Cut algorithm needs binary input image.\n");
        return;
        //exit(1);
    }

    // Mao, Fig.3: Step 1.
    intarray int_img_hor ;
    intarray int_img_ver ;
    compute_integral_image(in, int_img_hor, int_img_ver) ;

    // Mao, Fig.3: Step 2.
    rectarray blocks ;
    rectarray tree ;

    xycut(blocks, tree, int_img_hor, int_img_ver, tnx, tny, tcx, tcy) ;

    int cflen = blocks.length();

    makelike(image,in);
    fill(image,0x00ffffff);
    for(int i=0; i<cflen; i++){
        // Do a logical OR of zone index with R=1 to assign all
        // zones to the first column
        int color = (i+1)|(0x00010000);
        rectangle r = blocks[i];

        for(int x=r.x0; x<r.x1; x++){
            for(int y=r.y0; y<r.y1; y++){
                if(in(x,y) == 0)
                    image(x,y) = color ;
            }
        }
    }
    if(debug_xycut_intermediate){
        for (int i = 0; i < blocks.length(); i++) {
            fprintf(stderr,"blocks[%d] = %d\t%d\t%d\t%d\n", i, 
                    blocks[i].x0, blocks[i].y0, blocks[i].x1, blocks[i].y1);
        }
        write_image_packed("xycuts-cuts.png",image);
        write_image_packed("xycuts-inthor.png",int_img_hor);
        write_image_packed("xycuts-intver.png",int_img_ver);
    }
}

ISegmentPage *make_SegmentPageByXYCUTS() {
    return new SegmentPageByXYCUTS();
}

ISegmentPage *make_SegmentPageByXYCUTS(unsigned int itnx, 
        unsigned int itny, 
        unsigned int itcx, 
        unsigned int itcy) {
    return new SegmentPageByXYCUTS(itnx, itny, itcx, itcy);
}
